"""
Author: xiash$
Date: 2025/8/9$
Project: open-image$
Description: $
***
* _ooOoo_
* o8888888o
* 88" . "88
* (| -_- |)
*  O\ = /O
* ___/`---'\____
* .   ' \\| |// `.
* / \\||| : |||// \
* / _||||| -:- |||||- \
* | | \\\ - /// | |
* | \_| ''\---/'' | |
* \ .-\__ `-` ___/-. /
* ___`. .' /--.--\ `. . __
* ."" '< `.___\_<|>_/___.' >'"".
* | | : `- \`.;`\ _ /`;.`/ - ` : | |
* \ \ `-. \_ __\ /__ _/ .-` / /
* ======`-.____`-.___\_____/___.-`____.-'======
* `=---='
*          .............................................
*           佛曰：bug泛滥，我已瘫痪！
"""
import json
import math
import sys

import numpy as np
import random
from functools import partial
from typing import Union
from pathlib import Path
from concurrent.futures import ProcessPoolExecutor, as_completed
from rich.progress import Progress, MofNCompleteColumn, TimeElapsedColumn
from rich.panel import Panel
from rich import print
from collections import Counter
from itertools import product
from .open_image import open_image, OpenImage
from .enums import ImageFormat
from .open_shp import open_shp
from .utils import setup_logger
from .utils.converters import GeoConverter, ColorTypeConverter, PaletteConverter

logger = setup_logger(__name__)


def _get_save_directory_by_cycle(index: int, directories_dict: dict, ratios: list) -> tuple:
    """
    根据循环模式分配保存目录，确保数据均匀分布
    @param index: 图片序号（从1开始）
    @param directories: 目录字典，如 {'train':'训练集', 'val':'验证集', 'test':'测试集'}
    @param ratios: 对应比例，如 [0.7, 0.2, 0.1]，会转换为 [7, 2, 1] 的循环模式
    @return: 应该保存到的目录路径
    """
    if len(directories_dict) != len(ratios):
        raise ValueError("目录数量和比例数量必须相同!")
    list_directories = list(directories_dict.keys())
    if len(list_directories) == 1:
        save_dir = list_directories[0]
        return save_dir, directories_dict[save_dir]
    # 将比例转换为整数循环模式
    # 例如 [0.7, 0.2, 0.1] -> [7, 2, 1]
    # 找到最小比例单位
    min_ratio = min(ratios)
    cycle_pattern = [round(r / min_ratio) for r in ratios]
    # 构建循环序列
    cycle_sequence = []
    for dir_idx, count in enumerate(cycle_pattern):
        cycle_sequence.extend([list_directories[dir_idx]] * count)
    # 根据索引循环分配
    cycle_length = len(cycle_sequence)
    sequence_index = (index - 1) % cycle_length
    save_dir = cycle_sequence[sequence_index]
    return save_dir, directories_dict[save_dir]


class SemanticSampleBuilder:
    """
    语义分割样本制作
    """

    def __init__(
            self,
            input_images: list[str | Path] | str,  # 输入图像路径
            input_labels: list[str | Path] | str,  # 输入标签路径
            output_dir: str | Path,  # 输出目录
            attribute_field: str = 'class_id',  # shp属性字段
            mapping_field: dict = None,  # shp属性字段映射表
            window_size: int | tuple[int, int] = 512,  # 切片大小
            stride: int | tuple[int, int] = None,  # 步长
            max_samples: int = None,  # 最大切片数量
            band_indexes: int | list[int] = None,  # 波段索引
            palette: Union[list, tuple, dict, bytes, np.ndarray] = None,  # 颜色映射表
            split_val_ratio: float = 0,  # 验证集比例
            split_test_ratio: float = 0,  # 测试集比例
            save_blank_sample_prob: float = 0,  # 保存空白样本的概率
            save_file_name: str = None,  # 保存样本文件时的基础文件名
            image_prefix: str = None,  # 图像文件前缀
            label_prefix: str = None,  # 标签文件前缀
            image_suffix: str = None,  # 图像文件后缀
            label_suffix: str = None,  # 标签文件后缀
            input_mode: int = 0,  # 输入模式 0:文件路径列表,1:目录根据文件名匹配，2：目录根据地理范围匹配
            log_file: str = None,  # 日志输出到文件
            print_progress_rate: bool = True,  # 是否打印进度完成百分比
            num_workers: int = 4,  # 进程池线程数
    ):
        self.progress_bar = Progress(*Progress.get_default_columns(), MofNCompleteColumn(), TimeElapsedColumn())
        self.progress_task = self.progress_bar.add_task(f'数据准备中...', total=None)
        self.progress_bar.start()
        self.input_images = input_images
        self.input_labels = input_labels
        self.output_dir = Path(output_dir)
        self.attribute_field = attribute_field
        self.mapping_field = mapping_field
        self.window_size = window_size
        self.stride = stride or window_size
        self.max_samples = max_samples
        self.band_indexes = band_indexes
        self.palette = palette
        self.split_val_ratio = split_val_ratio
        self.split_test_ratio = split_test_ratio
        self.save_blank_sample_prob = save_blank_sample_prob
        self.save_file_name = save_file_name
        self.image_prefix = image_prefix
        self.label_prefix = label_prefix
        self.image_suffix = image_suffix
        self.label_suffix = label_suffix
        self.input_mode = input_mode
        self.logger = setup_logger(f'{__name__}.{self.__class__.__name__}', log_file=log_file)
        self.print_progress_rate = print_progress_rate
        self.save_dirs = {}
        self.split_ratios = [1 - self.split_val_ratio - self.split_test_ratio]
        self.input_data = []
        self.progress_total = 0
        self.progress_completed = 0
        self.executor = ProcessPoolExecutor(max_workers=num_workers)
        self.prepare_data()
        self.init_save_dir()
        self.init_progress()
        self.update_progress(description='数据准备完成')

    def close(self):
        self.executor.shutdown()
        self.progress_bar.stop()

    def init_progress(self):
        """
        初始化进度条
        """
        if isinstance(self.stride, int):
            self.stride = (self.stride, self.stride)
        stride_x, stride_y = self.stride[0], self.stride[-1]
        for image_path, label_path in self.input_data:
            image = open_image(str(image_path))
            label = open_shp(label_path) if label_path.suffix == '.shp' else open_image(label_path)
            if hasattr(label, 'geo_transform') and label.geo_transform == (0, 1, 0, 0, 0, 1):
                pixel_bounds = [0, 0, image.width, image.height]
            else:
                if not OpenImage.is_same_projection(image.geo_projection, label.geo_projection):
                    label.reproject(image.geo_projection)
                geo_bounds_list = OpenImage.intersection_bounds([image.geo_bounds, label.geo_bounds])
                if not geo_bounds_list:
                    raise ValueError(f'图像[{image.file_name}]和标签[{label.file_name}]地理范围不匹配！')
                pixel_bounds = GeoConverter.geo_bounds2pixel_bounds(geo_bounds_list, image.geo_transform)

            crop_width = pixel_bounds[2] - pixel_bounds[0]
            crop_height = pixel_bounds[3] - pixel_bounds[1]
            self.progress_total += math.ceil(crop_width / stride_x) * math.ceil(crop_height / stride_y)
            image.close()
            label.close()
        self.progress_total += len(self.input_data) + 1
        self.progress_bar.update(self.progress_task, total=self.progress_total)

    def init_save_dir(self):
        assert self.split_ratios[0] > 0, '验证集和测试集比例之和不能大于等于1!'
        self.save_dirs[self.output_dir / 'train'] = '训练集'
        if self.split_val_ratio:
            self.save_dirs[self.output_dir / 'val'] = '验证集'
            self.split_ratios.append(self.split_val_ratio)
        if self.split_test_ratio:
            self.save_dirs[self.output_dir / 'test'] = '测试集'
            self.split_ratios.append(self.split_test_ratio)
        for save_dir in self.save_dirs:
            image_save_dir = save_dir / 'images'
            label_save_dir = save_dir / 'labels'
            image_save_dir.mkdir(parents=True, exist_ok=True)
            label_save_dir.mkdir(parents=True, exist_ok=True)

    def update_progress(self, advance=1., description=None):
        self.progress_completed += advance
        self.progress_bar.update(self.progress_task, completed=self.progress_completed, description=description)
        if self.print_progress_rate:
            if self.progress_completed % 100 == 1:
                self.logger.info(f'[PROGRESS][{self.progress_completed * 100 / self.progress_total:.2f}%]{description}')
            elif self.progress_completed == self.progress_total:
                self.logger.info(f'[PROGRESS][100%]{description}')

    def build(self):
        """
        制作样本
        """
        save_index = 0
        sample_counter = Counter()
        try:
            for image_path, label_path in self.input_data:
                image, label = self.read_one_data(image_path, label_path, self.attribute_field, self.mapping_field)
                self.update_progress(
                    description=(desc := f'[图像：{image.file_name}][标签：{label.file_name}]样本切片中...')
                )
                image_tiles = image.sliding_window_tiles(self.window_size, self.stride, band_indexes=self.band_indexes)
                label_tiles = label.sliding_window_tiles(self.window_size, self.stride)
                futures = {}
                save_file_name = self.save_file_name or image.file_name
                for image_tile, label_tile in zip(image_tiles, label_tiles):
                    save_index += 1
                    image_array, label_array = image_tile[0], label_tile[0]
                    if np.all(label_array == 0):
                        sample_counter['空白标签样本数量'] += 1
                        if not random.choices(
                                [True, False],
                                [self.save_blank_sample_prob, 1 - self.save_blank_sample_prob]
                        )[0]:
                            self.update_progress(description=desc)
                            sample_counter['删除空白标签样本数量'] += 1
                            continue
                    save_dir, dir_name = _get_save_directory_by_cycle(save_index, self.save_dirs, self.split_ratios)
                    image_save_path = save_dir / f'images/{save_file_name}_{save_index}.png'
                    label_save_path = save_dir / f'labels/{save_file_name}_{save_index}.png'
                    func = partial(OpenImage.write_array, backend='pil')
                    futures[self.executor.submit(func, image_save_path, image_array, mode='RGB')] = dir_name
                    futures[
                        self.executor.submit(func, label_save_path, label_array, mode='L', palette=self.palette)
                    ] = dir_name
                for i, future in enumerate(as_completed(futures), start=1):
                    future.result()
                    if i % 2 == 0:
                        self.update_progress(description=desc)
                        sample_counter[f'{futures.pop(future)}数量'] += 1
                        sample_counter['保留样本数量'] += 1
                        if sample_counter['保留样本数量'] == self.max_samples:
                            self.update_progress(
                                advance=self.progress_total - self.progress_completed,
                                description=f'已生成设置的最大样本数量:{self.max_samples}，程序提前结束'
                            )
                            sys.exit(1)

                image.close()
                label.close()

            return sample_counter
        except Exception as e:
            self.logger.exception(e)
        finally:
            sample_counter['总切片数量'] = save_index
            SemanticSampleBuilder.print_sample_info(
                sample_counter, tile=f'{len(self.input_data)}组源图和标签生成总样本量统计'
            )
            SemanticSampleBuilder.save_sample_info(
                self.output_dir, sample_counter, self.window_size, self.stride, self.palette
            )
            self.update_progress(advance=self.progress_total - self.progress_completed, description='样本切片完成！')
            self.close()

    def prepare_data(self):
        """
        读取所有图像和标签
        """
        all_image_format = ImageFormat.all_formats()
        input_images = SemanticSampleBuilder.input_to_list_files(self.input_images, all_image_format)
        input_labels = SemanticSampleBuilder.input_to_list_files(self.input_labels, ['.shp'] + all_image_format)
        if self.input_mode == 0:
            self.input_data = list(zip(input_images, input_labels))
        else:
            for image_path, label_path in product(input_images, input_labels):
                if SemanticSampleBuilder.image_matched_label(image_path, label_path, self.image_prefix,
                                                             self.label_prefix,
                                                             self.image_suffix, self.label_suffix, self.input_mode):
                    self.input_data.append((image_path, label_path))

    @staticmethod
    def input_to_list_files(input_data: list[str | Path] | str, pattern: str | list[str]) -> list[Path]:
        if isinstance(input_data, str):
            input_data = Path(input_data)
            if input_data.is_file():
                input_data = [input_data]
            elif input_data.is_dir():
                input_data = [file for file in input_data.rglob('*') if file.suffix in pattern]
            else:
                raise ValueError(f'输入路径错误：{input_data}')
        else:
            input_data = [Path(file) for file in input_data if Path(file).suffix in pattern]
        return input_data

    @staticmethod
    def image_matched_label(
            image_path: str | Path,
            label_path: str | Path,
            image_prefix: str = None,
            label_prefix: str = None,
            image_suffix: str = None,
            label_suffix: str = None,
            matching_mode: int = 1
    ):
        """
        判断图像和标签是否匹配
        @param image_path: 图像路径
        @param label_path: 标签路径
        @param image_prefix: 图像文件前缀
        @param label_prefix: 标签文件前缀
        @param image_suffix: 图像文件后缀
        @param label_suffix: 标签文件后缀
        @param matching_mode: 匹配模式,1:根据文件名，2：根据地理范围
        @return: 是否匹配
        """
        image_path, label_path = Path(image_path), Path(label_path)
        if matching_mode == 1:
            image_name, label_name = image_path.name, label_path.name
            image_suffix = image_suffix or image_path.suffix
            label_suffix = label_suffix or label_path.suffix
            image_prefix = image_prefix or ''
            label_prefix = label_prefix or ''
            return image_name.removeprefix(image_prefix).removesuffix(image_suffix) == label_name \
                .removeprefix(label_prefix).removesuffix(label_suffix)
        image = open_image(str(image_path))
        image_bounds = image.geo_bounds
        if ImageFormat.is_image_format(label_path):
            label = open_image(str(label_path))
        elif label_path.suffix == '.shp':
            label = open_shp(str(label_path))
        else:
            return False
        if not OpenImage.is_same_projection(image.geo_projection, label.geo_projection):
            label.reproject(image.geo_projection)
        label_bounds = label.geo_bounds
        image.close()
        label.close()
        intersection_bounds = OpenImage.intersection_bounds([image_bounds, label_bounds])
        if not intersection_bounds:
            return False
        # intersection_w = (intersection_bounds[2] - intersection_bounds[0])
        # intersection_h = abs(intersection_bounds[1] - intersection_bounds[3])
        # image_w = (image_bounds[2] - image_bounds[0])
        # image_h = abs(image_bounds[1] - image_bounds[3])
        # if intersection_w * intersection_h / (image_w * image_h) < 0.5:
        #     return False
        return True

    @staticmethod
    def read_one_data(
            image_path: str | Path,
            label_path: str | Path,
            attribute_field: str = None,
            mapping_field: dict = None
    ) -> tuple[OpenImage, OpenImage]:
        image = open_image(str(image_path))
        if Path(label_path).suffix == '.shp':
            with open_shp(label_path) as shp:
                if not OpenImage.is_same_projection(image.geo_projection, shp.geo_projection):
                    logger.warning(f'标签[{shp.file_name}]和图像[{image.file_name}]的投影不匹配，将重新投影')
                    shp.reproject(image.geo_projection)
                geo_bounds = OpenImage.intersection_bounds([image.geo_bounds, shp.geo_bounds])
                image = image.crop_by_geo_bounds(geo_bounds)
                img_data = GeoConverter.shp_to_raster(
                    shp.data_source,
                    out_width=image.width,
                    out_height=image.height,
                    output_path='',
                    attribute_field=attribute_field,
                    mapping_field=mapping_field,
                    geo_transform=image.geo_transform,
                    projection=image.geo_projection,
                    data_type=image.img_data.GetRasterBand(1).DataType
                )
                label = open_image(shp.file_path, 'w')
                label.set_img_data(img_data)
        else:
            label = open_image(str(label_path))
            if label.geo_transform == (0, 1, 0, 0, 0, 1):
                label = label.set_backend_image(label.resize(image.size))
                label.set_geo_transform(image.geo_transform).set_geo_projection(image.geo_projection)
            else:
                image, label = OpenImage.intersection_image([image, label])
                label = label.set_backend_image(label.resize(image.size))
        return image, label

    @staticmethod
    def print_sample_info(list_sample_info: list[dict] | dict, tile: str = None):
        """
        打印样本信息

        @param list_sample_info: 样本信息列表
        @param tile: 标题
        @return: None
        """
        if isinstance(list_sample_info, dict):
            list_sample_info = [list_sample_info]
        print_info = Counter()
        for sample_info in list_sample_info:
            print_info += Counter(sample_info)
        print(Panel(
            '\n'.join([f'{key}: {value}' for key, value in print_info.items()]),
            title=tile,
            border_style="green"
        ))

    @staticmethod
    def build_color_label(
            ori_label_path: str,
            color_label_path: str,
            palette: Union[list, tuple, dict, bytes, np.ndarray],
            rgb_reversed_mapping: dict[tuple, int] = None,
            sieve_size: int = None
    ):
        """
        将原始标签转为伪彩色标签

        @param ori_label_path: 原始标签路径
        @param color_label_path: 伪彩标签路径
        @param palette: 颜色标签的调色板，支持pil，rasterio等风格
        @param rgb_reversed_mapping: RGB颜色反向映射，当原始标签图像是rgb三波段时设置，例如：{(255,0,0):1}
        @param sieve_size: 过滤目标像素尺寸
        @return: None
        """
        from PIL import Image
        Image.MAX_IMAGE_PIXELS = None
        with (
            open_image(ori_label_path, backend='pil') as ori_label,
            open_image(color_label_path, mode='L', backend='pil') as color_label
        ):
            match ori_label.channel:
                case 1:
                    color_label_array = ori_label.img_array
                case i if i >= 3:
                    assert rgb_reversed_mapping is not None, '请设置RGB颜色反向映射！'
                    ori_label_array = ori_label.img_array[..., :3]
                    color_label_array = np.zeros((ori_label.height, ori_label.width), dtype=np.uint8)
                    for (r, g, b), class_id in rgb_reversed_mapping.items():
                        color_label_array[
                            (ori_label_array[..., 0] == r) & (ori_label_array[..., 1] == g) & (
                                    ori_label_array[..., 2] == b)
                            ] = class_id
                case _:
                    raise ValueError(f'{ori_label_path} 原始标签图像格式有误，必须是单波段或大于等于三波段！')
            if sieve_size is not None:
                from rasterio.features import sieve
                color_label_array = sieve(color_label_array, sieve_size, connectivity=8)
            color_label.set_img_array(color_label_array).set_palette(palette).save()

    @staticmethod
    def batch_build_color_label(
            ori_label_dir: str,
            color_label_dir: str,
            palette: Union[list, tuple, dict, bytes, np.ndarray],
            rgb_reversed_mapping: dict[tuple, int] = None,
            sieve_size: int = None,
    ):
        """
        批量将整个目录的原始标签转为伪彩色标签

        @param ori_label_dir: 原始标签目录
        @param color_label_dir: 伪彩标签保存目录
        @param palette: 颜色表
        @param rgb_reversed_mapping: rgb->class_id 映射表，标签为rgb图像时使用，例如：{(255,0,0):1}
        @param sieve_size: 过滤小目标像素面积阈值
        """
        with (
            ProcessPoolExecutor(4) as executor,
            Progress(*Progress.get_default_columns(), MofNCompleteColumn(), TimeElapsedColumn()) as progress
        ):
            p_bar = progress.add_task('[green]批量转伪彩标签图像...', total=None)
            ori_label_dir, color_label_dir = Path(ori_label_dir), Path(color_label_dir)
            futures = []
            for file in ori_label_dir.iterdir():
                if not ImageFormat.is_image_format(file.suffix):
                    continue
                ori_label_path = str(file)
                color_label_path = str(color_label_dir / f'{file.stem}.png')
                futures.append(executor.submit(
                    SemanticSampleBuilder.build_color_label,
                    ori_label_path,
                    color_label_path,
                    palette,
                    rgb_reversed_mapping,
                    sieve_size
                ))
            progress.update(p_bar, total=len(futures))
            for future in as_completed(futures):
                future.result()
                progress.advance(p_bar)
            progress.update(p_bar, description='[purple]批量转伪彩标签图像完成！')

    @staticmethod
    def build_visualizer_label(
            image_path: str,
            label_path: str,
            palette: Union[list, tuple, dict, bytes, np.ndarray],
            visualizer_label_path: str = None,
            contour_width: int = 2,
            class_names: dict[int, str] = None,
            font_size: int = 8,
            font_color: str | tuple[int, int, int] = '白色',
            font_path: str = 'simhei.ttf',
            show=False
    ):
        """
        将语义分割标签转图像轮廓画到原图上

        @param image_path: 原始图像路径
        @param label_path: 标签路径
        @param palette: 调色板,支持pil，rasterio等多种风格，且支持常见中英文颜色名称如"red","yellow","blue"等
        @param visualizer_label_path: 可视化标签路径，不为空时会保存结果
        @param contour_width: 轮廓线宽度
        @param class_names: 所有类别名称，格式为字典如 {1:'建筑',2:'道路'}
        @param font_size: 字体大小
        @param font_color: 字体颜色
        @param font_path: 字体文件路径
        @param show: 是否显示图片
        @return: None
        """
        from PIL import Image, ImageDraw, ImageFont
        from skimage.measure import find_contours

        # 加载原图和mask
        original_image = Image.open(image_path)
        mask = Image.open(label_path)
        draw = ImageDraw.Draw(original_image)
        # 定义类别名称（支持中文）和颜色（RGB格式）
        palette: dict = PaletteConverter.convert('rasterio', palette)
        try:
            font = ImageFont.truetype(font_path, font_size)
        except Exception:
            logger.warning(
                f'{font_path} 字体文件加载失败！将使用默认字体，不支持中文。如需中文支持，请自行下载字体文件并指定字体文件路径！',
                markup='yellow')
            font = None
        font_color = ColorTypeConverter.to_rgb_tuple(font_color)
        # 将mask转换为NumPy数组
        mask_array = np.array(mask)
        0 in palette and palette.pop(0)
        for class_id, class_color in palette.items():
            # 创建当前类别的二值mask数组（0/1）
            binary_mask = (mask_array == class_id).astype(bool)  # find_contours需要float或bool
            # 使用scikit-image提取精确轮廓（level=0.5表示边界）
            contours = find_contours(binary_mask, level=0.5)
            # 绘制每个轮廓（contours是list of arrays，每个array是(y, x)坐标）
            for contour in contours:
                # 转换为整数坐标并绘制多段线（精确轮廓）
                contour = contour.astype(int)
                # PIL的line需要(x, y)列表，所以翻转轴并转换为list of tuples
                contour_list = [(x, y) for y, x in contour]  # 注意：find_contours返回(row, col)即(y, x)
                draw.line(contour_list, fill=class_color, width=contour_width)
                # 计算这个轮廓的质心（使用轮廓点的均值）
                if class_names is not None:
                    cy, cx = np.mean(contour, axis=0).astype(int)  # (y, x)
                    text = class_names.get(class_id, f'class_{class_id}')
                    # 获取文本边界框以居中
                    text_bbox = draw.textbbox((cx, cy), text, font=font)
                    text_width = text_bbox[2] - text_bbox[0]
                    text_height = text_bbox[3] - text_bbox[1]
                    draw.text((cx - text_width // 2, cy - text_height // 2), text, fill=font_color, font=font,
                              font_size=font_size)

        show and original_image.show()
        if visualizer_label_path is not None:
            Path(visualizer_label_path).parent.resolve().mkdir(parents=True, exist_ok=True)
            # 确保图像不是P模式，如果是则转换为RGB模式再保存为JPEG
            original_image.convert('RGB').save(visualizer_label_path)
        return original_image

    @staticmethod
    def batch_build_visualizer_label(
            image_dir: str,
            label_dir: str,
            visualizer_label_dir: str,
            palette: Union[list, tuple, dict, bytes, np.ndarray],
            contour_width: int = 2,
            class_names: dict[int, str] = None,
            font_size: int = 8,
            font_color: str | tuple[int, int, int] = '白色',
            font_path: str = 'simhei.ttf',
            prefixes: tuple[str, str] = None,
            suffixes: tuple[str, str] = None,

    ):
        """
        批量将目录下的所有语义分割标签转图像轮廓画到原图上

        @param image_dir: 原始图像文件夹路径
        @param label_dir: 语义分割标签文件夹路径
        @param visualizer_label_dir: 可视化标签文件夹路径
        @param palette: 调色板,支持pil，rasterio等多种风格，且支持常见中英文颜色名称如"red","yellow","blue"等
        @param contour_width: 轮廓线宽度
        @param class_names: 所有类别名称，格式为字典如 {1:'建筑',2:'道路'}
        @param font_size: 字体大小
        @param font_color: 字体颜色
        @param font_path: 字体路径
        @param prefixes: 文件名前缀，格式如：('img_','label_')
        @param suffixes: 文件名后缀，格式如：('_img.tif','_label.png')
        """
        with (
            ProcessPoolExecutor(4) as executor,
            Progress(*Progress.get_default_columns(), MofNCompleteColumn(), TimeElapsedColumn()) as progress
        ):
            p_bar = progress.add_task('[green]批量可视化标签图像...', total=None)
            image_dir, label_dir, visualizer_label_dir = Path(image_dir), Path(label_dir), Path(visualizer_label_dir)
            futures = []
            run_func = partial(SemanticSampleBuilder.build_visualizer_label,
                               palette=palette,
                               contour_width=contour_width,
                               class_names=class_names,
                               font_size=font_size,
                               font_color=font_color,
                               font_path=font_path)
            for file in image_dir.iterdir():
                if not ImageFormat.is_image_format(file.suffix):
                    continue
                image_path = str(file)
                label_file_name = file.name
                if prefixes is not None:
                    image_prefix, label_prefix = prefixes
                    label_file_name = label_file_name.replace(image_prefix, label_prefix)
                if suffixes is not None:
                    image_suffix, label_suffix = suffixes
                    label_file_name = label_file_name.replace(image_suffix, label_suffix)
                label_path = str(label_dir / label_file_name)
                visualizer_label_path = str(visualizer_label_dir / label_file_name)
                futures.append(executor.submit(
                    run_func,
                    image_path=image_path,
                    label_path=label_path,
                    visualizer_label_path=visualizer_label_path
                ))
            progress.update(p_bar, total=len(futures))
            for future in as_completed(futures):
                future.result()
                progress.advance(p_bar)
            progress.update(p_bar, description='[purple]批量可视化标签图像完成！')

    @staticmethod
    def save_sample_info(output_dir, total_sample_counter, window_size, stride, palette=None):
        sample_info = dict(total_sample_counter)
        sample_info.update({
            '切片大小': window_size,
            '切割步长': stride or window_size,
        })
        palette and sample_info.update({'标签颜色映射表': PaletteConverter.convert('rasterio', palette)})
        save_path = str(Path(output_dir) / 'sample_info.json')
        with open(save_path, 'w', encoding='utf-8') as f:
            json.dump(sample_info, f, ensure_ascii=False, indent=4)


class ChangedSampleBuilder(SemanticSampleBuilder):
    @staticmethod
    def build_change_sample_by_open_image(
            output_dir: str,
            image1: OpenImage,
            image2: OpenImage,
            mask: OpenImage,
            window_size: int | tuple[int, int] = 512,
            stride: int | tuple[int, int] = None,
            max_samples: int = None,
            band_indexes: int | list[int] = None,
            palette: Union[list, tuple, dict, bytes, np.ndarray] = None,
            split_val_ratio: float = None,
            split_test_ratio: float = None,
            save_blank_sample_prob: float = 0,
            save_file_name=None
    ) -> Counter:
        """
        基于OpenImage构建变化检测样本
        @param output_dir: 输出根目录
        @param image1: 前期图像路径
        @param image2: 后期图像路径
        @param mask: 标签图像路径
        @param window_size: 切片大小，(width, height)
        @param stride: 步长，(x_step, y_step)，默认与window_size相同
        @param max_samples: 最大切片数量，默认为None（不限制）
        @param band_indexes: 选择的图像波段索引，默认为None（使用所有波段）
        @param palette: 颜色映射表
        @param split_val_ratio: 验证集比例
        @param split_test_ratio: 测试集比例
        @param save_blank_sample_prob: 保存空白样本的概率
        @param save_file_name: 保存样本文件时的基础文件名，默认会以原图名称命名
        @return: 样本信息字典
        """
        output_dir = Path(output_dir)
        image1_save_dirs = [output_dir / 'train/images1']
        image2_save_dirs = [output_dir / 'train/images2']
        label_save_dirs = [output_dir / 'train/labels']
        ratios = [1, ]
        if split_val_ratio:
            image1_save_dirs.append(output_dir / 'val/images1')
            image2_save_dirs.append(output_dir / 'val/images2')
            label_save_dirs.append(output_dir / 'val/labels')
            ratios.append(split_val_ratio)
        if split_test_ratio:
            image1_save_dirs.append(output_dir / 'test/images1')
            image2_save_dirs.append(output_dir / 'test/images2')
            label_save_dirs.append(output_dir / 'test/labels')
            ratios.append(split_test_ratio)
        ratios[0] = 1 - sum(ratios[1:])
        assert ratios[0] > 0, '验证集和测试集比例之和不能大于等于1!'
        if image2.geo_transform == (0, 1, 0, 0, 0, 1):
            image2.size != image1.size and image2.resize(image1.size)
            image2.set_geo_projection(image1.geo_projection).set_geo_transform(image1.geo_transform)
        if mask.geo_transform == (0, 1, 0, 0, 0, 1):
            mask.size != image1.size and mask.resize(image1.size)
            mask.set_geo_projection(image1.geo_projection).set_geo_transform(image1.geo_transform)
        image1, image2, mask = OpenImage.intersection_image([image1, image2, mask])
        image1.size != image2.size and image2.resize(image1.size)
        image1.size != mask.size and mask.resize(image1.size)
        image1_tiles = image1.sliding_window_tiles(window_size, stride, max_samples, band_indexes)
        image2_tiles = image2.sliding_window_tiles(window_size, stride, max_samples, band_indexes)
        mask_tiles = mask.sliding_window_tiles(window_size, stride, max_samples)
        for save_dirs in zip(image1_save_dirs, image2_save_dirs, label_save_dirs):
            for save_dir in save_dirs:
                save_dir.mkdir(parents=True, exist_ok=True)
        futures = []
        with (
            ProcessPoolExecutor(8) as executor,
            Progress(*Progress.get_default_columns(), MofNCompleteColumn(), TimeElapsedColumn()) as progress
        ):
            title = f'(前期：{image1.file_name}{image1.img_suffix}，后期：{image2.file_name}{image2.img_suffix}，标签：{mask.file_name}{mask.img_suffix})'
            task = progress.add_task(f'[green]{title}：变化检测样本切片生成中...', total=None)
            sample_counter = Counter()
            if not save_file_name:
                save_file_name = image1.file_name
            for i, (image1_tile, image2_tile, mask_tile) in enumerate(zip(image1_tiles, image2_tiles, mask_tiles),
                                                                      start=1):
                image1_array, image2_array, mask_array = image1_tile[0], image2_tile[0], mask_tile[0]
                sample_counter['切片数量'] += 1
                if np.all(mask_array == 0):
                    sample_counter['空白标签样本数量'] += 1
                    if not random.choices([True, False], [save_blank_sample_prob, 1 - save_blank_sample_prob])[0]:
                        progress.advance(task)
                        sample_counter['删除空白标签样本数量'] += 1
                        continue
                image1_save_dir = _get_save_directory_by_cycle(i, image1_save_dirs, ratios)
                image1_save_path = Path(image1_save_dir) / f'{save_file_name}_{i}.png'
                save_dir_index = image1_save_dirs.index(image1_save_dir)
                image2_save_dir = image2_save_dirs[save_dir_index]
                image2_save_path = Path(image2_save_dir) / f'{save_file_name}_{i}.png'
                label_save_dir = label_save_dirs[save_dir_index]
                label_save_path = Path(label_save_dir) / f'{save_file_name}_{i}.png'
                func = partial(OpenImage.write_array, backend='pil')
                futures.append(executor.submit(
                    func,
                    image1_save_path,
                    image1_array,
                    mode='RGB'
                ))
                futures.append(executor.submit(
                    func,
                    image2_save_path,
                    image2_array,
                    mode='RGB'
                ))
                futures.append(executor.submit(
                    func,
                    label_save_path,
                    mask_array,
                    mode='L',
                    palette=palette
                ))
                if save_dir_index == 0:
                    sample_counter['训练集数量'] += 1
                if save_dir_index == 1:
                    sample_counter['验证集数量'] += 1
                if save_dir_index == 2:
                    sample_counter['测试集数量'] += 1
            progress.update(task, total=i)
            for future in as_completed(futures):
                future.result()
                progress.update(task, advance=1 / 3)
                sample_counter['保留样本数量'] += 1
            progress.update(task, description=f'[purple]{title}: 变化检测样本制作完成！', completed=i)
            sample_counter['保留样本数量'] //= 3
            ChangedSampleBuilder.print_sample_info(sample_counter, tile=title)
            ChangedSampleBuilder.save_sample_info(output_dir, sample_counter, window_size, stride, palette)
            image1.close()
            image2.close()
            mask.close()
            return sample_counter

    @staticmethod
    def build_change_sample(
            output_dir: str,
            image1_path: str,
            image2_path: str,
            label_path: str,
            attribute_field: str = None,
            mapping_field: dict = None,
            window_size: int | tuple[int, int] = 512,
            stride: int | tuple[int, int] = None,
            max_samples: int = None,
            band_indexes: int | list[int] = None,
            palette: Union[list, tuple, dict, bytes, np.ndarray] = None,
            split_val_ratio: float = None,
            split_test_ratio: float = None,
            save_blank_sample_prob: float = 0,
            save_file_name: str = None
    ) -> Counter:
        """
        基于栅格或矢量标签数据构建变化检测样本
        @param output_dir: 输出根目录
        @param image1_path: 前期图像路径
        @param image2_path: 后期图像路径
        @param label_path: 标签图像路径或者shp文件路径
        @param attribute_field: 用于栅格化的属性字段名
        @param mapping_field: 字段映射，例如：{'建筑':1,'道路':2}
        @param window_size: 切片大小，(width, height)
        @param stride: 步长，(x_step, y_step)，默认与window_size相同
        @param max_samples: 最大切片数量，默认为None（不限制）
        @param band_indexes: 选择的图像波段索引，默认为None（使用所有波段）
        @param palette: 调色板，用于标签图像着色
        @param split_val_ratio: 验证集比例，默认为None（不划分）
        @param split_test_ratio: 测试集比例，默认为None（不划分）
        @param save_blank_sample_prob: 保留空白标签样本的概率，0到1之间
        @param save_file_name: 保存样本文件时的基础文件名，默认会以原图名称命名
        @return: 样本信息字典
        """
        image1, image2 = open_image(image1_path), open_image(image2_path)
        mask = ChangedSampleBuilder.read_one_data(label_path, attribute_field, image1, mapping_field)
        return ChangedSampleBuilder.build_change_sample_by_open_image(
            output_dir,
            image1,
            image2,
            mask,
            window_size,
            stride,
            max_samples,
            band_indexes,
            palette,
            split_val_ratio,
            split_test_ratio,
            save_blank_sample_prob,
            save_file_name
        )

    @staticmethod
    def batch_build_change_sample(
            output_dir: str,
            image1_dir: str,
            image2_dir: str,
            label_dir: str,
            attribute_field: str = None,
            mapping_field: dict = None,
            window_size: int | tuple[int, int] = 512,
            stride: int | tuple[int, int] = None,
            max_samples: int = None,
            band_indexes: int | list[int] = None,
            palette: Union[list, tuple, dict, bytes, np.ndarray] = None,
            split_val_ratio: float = None,
            split_test_ratio: float = None,
            save_blank_sample_prob=0,
            save_file_name=None,
            prefixes: tuple[str, str, str] | list[str] = ('', '', ''),
            suffixes: tuple[str, str, str] | list[str] = None,
            matching_mode: int = 0
    ) -> Counter:
        """
        顺序扫描目录，批量制作变化检测样本
        @param output_dir:输出目录
        @param image1_dir:前期图像所在目录
        @param image2_dir:后期图像所在目录
        @param label_dir:标签图像或者shp所在目录
        @param attribute_field: 矢量转栅格的属性字段
        @param mapping_field: 属性字段映射，如 {'建筑':1,'道路':2}
        @param window_size:滑动窗口大小(width, height)
        @param stride:滑动步长(x_step, y_step)
        @param max_samples:最大生成样本数量
        @param band_indexes:源图像波段索引列表
        @param palette:颜色映射表
        @param split_val_ratio:验证集比例
        @param split_test_ratio:测试集比例
        @param save_blank_sample_prob: 保留空白标签样本的概率，0到1之间
        @param save_file_name: 保存样本文件时的基础文件名，默认会以原图名称命名
        @param prefixes: 源图像和标签文件名的前缀
        @param suffixes: 原图像和标签文件名的后缀
        @param matching_mode: 原图像和标签匹配的模式，0：根据文件名匹配，1：根据图像和标签的地理范围去匹配
        @return:
        """
        total_sample_counter = Counter()
        save_index = 0
        logger.info(f"扫描目录匹配图像和标签文件...", markup='green')
        for image1_file, image2_file, label_file in product(
                Path(image1_dir).iterdir(),
                Path(image2_dir).iterdir(),
                Path(label_dir).iterdir()):  # type: Path
            if ImageFormat.is_image_format(image1_file.suffix) and ImageFormat.is_image_format(image2_file.suffix) and (
                    ImageFormat.is_image_format(label_file.suffix) or label_file.suffix == '.shp'):
                image1_prefix, image2_prefix, label_prefix = prefixes
                image1_suffix, image2_suffix, label_suffix = suffixes if suffixes else (
                    image1_file.suffix, image2_file.suffix, label_file.suffix)
                if ChangedSampleBuilder.image_matched_label(image1_file, label_file, image1_prefix, label_prefix,
                                                            image1_suffix,
                                                            label_suffix,
                                                            matching_mode) and ChangedSampleBuilder.image_matched_label(
                    image2_file, label_file,
                    image2_prefix, label_prefix,
                    image2_suffix, label_suffix,
                    matching_mode):
                    remaining_sample = max_samples - total_sample_counter['保留样本数量'] if max_samples else None
                    if remaining_sample == 0:
                        break
                    save_index += 1
                    base_name = f'{save_file_name}_{save_index}' if save_file_name else image1_file.name.removeprefix(
                        image1_prefix).removesuffix(image1_suffix)
                    total_sample_counter += ChangedSampleBuilder.build_change_sample(
                        output_dir,
                        str(image1_file),
                        str(image2_file),
                        str(label_file),
                        attribute_field=attribute_field,
                        mapping_field=mapping_field,
                        window_size=window_size,
                        stride=stride,
                        max_samples=remaining_sample,
                        band_indexes=band_indexes,
                        palette=palette,
                        split_val_ratio=split_val_ratio,
                        split_test_ratio=split_test_ratio,
                        save_blank_sample_prob=save_blank_sample_prob,
                        save_file_name=base_name
                    )
        ChangedSampleBuilder.print_sample_info(total_sample_counter, tile=f'{save_index}组源图和标签生成总样本量统计')
        ChangedSampleBuilder.save_sample_info(output_dir, total_sample_counter, window_size, stride, palette)
        return total_sample_counter
